# -*- encoding: utf-8 -*-
'''

Openvino inference model, load pre-trained model and initialize NCS2 for inference.

@File    :   model.py
@Time    :   2021/12/04 15:20:44
@Author  :   Yangshuo He
@Contact :   sugarhe58@gmail.com
'''

import logging as log
import os
import sys

import cv2
import numpy as np
from openvino.inference_engine import IECore

class OpenvinoModel():
    def __init__(self, model_path) -> None:
        self.model_path = model_path
        self.ie = IECore()
        
    def readModel(self):
        """
        Read model from .xml file.
        Output:
            True - read success
            False - read failed
        """
        self.net = self.ie.read_network(model=self.model_path)

        if len(self.net.input_info) != 1:
            print("The sample supports only single input topologies")
            return False

        if len(self.net.outputs) != 1 and not ("boxes" in self.net.outputs or "labels" in self.net.outputs):
            print("The sample supports models with 1 output or with 2 with the names 'boxes' and 'labels'")
            return False
        
        return True

    def configureModel(self):
        """
        Configure input & output precision
        """
        print("Configuring input and output blobs")
        # Get name of input blob
        self.input_blob = next(iter(self.net.input_info))

        # Set input and output precision manually
        self.net.input_info[self.input_blob].precision = "U8"

        self.output_blob = next(iter(self.net.outputs))
        self.net.outputs[self.output_blob].precision = "FP32"

    def loadModel(self):
        """
        Load model to NCS2.
        """
        print("Loading the model to the plugin")
        self.exec_net = self.ie.load_network(network=self.net, device_name="MYRIAD")

    def processInput(self, image):
        """
        Resize image to required input size, re-arrange layout of input tensor.
        Input:
            image - input image, layout (HWC)
        Output:
            image - pre-processed image, layout (NCHW)
        """
        # Get network input shape
        _, _, net_h, net_w = self.net.input_info[self.input_blob].input_data.shape

        # resize image
        if image.shape[:-1] != (net_h, net_w):
            image = cv2.resize(image, (net_w, net_h))

        # Change data layout from HWC to CHW
        image = image.transpose((2, 0, 1))
        # Add N dimension to transform to NCHW
        image = np.expand_dims(image, axis=0)

        return image

    def inference(self, image):
        """
        Inference on NCS2.
        Input:
            image - resized image, for inference
        Output:
            detection - detection result for the first object, format [image_id, label, conf, x_min, y_min, x_max, y_max], the last five value range [0, 1]
        """

        result = self.exec_net.infer(inputs={self.input_blob: image})

        result = result[self.output_blob]
        # Change a shape of a numpy.ndarray with results ([1, 1, N, 7]) to get another one ([N, 7]), where N is the number of detected bounding boxes
        # Only take the first detection to accelerate
        detection = result.reshape(-1, 7)[0, :]
        
        return detection
            
if __name__ == '__main__':

    # log.basicConfig(format="[ %(levelname)s ] %(message)s", level=log.INFO, stream=sys.stdout)

    model_dir = "/home/pi/openvino/open_model_zoo/models/intel/"
    model_name = "person-detection-0200"
    model_acc = "FP16"

    # (.xml and .bin files) or (.onnx file)
    model_path = "{}/{}/{}/{}.xml".format(model_dir, model_name, model_acc, model_name)

    original_image = cv2.imread("/home/pi/Pictures/supervise/test1.jpg")

    openvino_model  = OpenvinoModel(model_path)
    read_fb = openvino_model.readModel()
    if read_fb:
        print("Reading the network: {}".format(model_name))
    else:
        print("Read network failed!")

    openvino_model.configureModel()
    openvino_model.loadModel()
    image_process = openvino_model.processInput(original_image)
    detection = openvino_model.inference(image_process)

    _, class_id, confidence, xmin, ymin, xmax, ymax = detection

    output_image = original_image.copy()

    h, w, _ = output_image.shape


    if confidence > 0.5:
        xmin = int(xmin * w)
        ymin = int(ymin * h)
        xmax = int(xmax * w)
        ymax = int(ymax * h)

        cv2.rectangle(output_image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)

    output_filename = "out_{}_{}.bmp".format(model_name, model_acc)
    cv2.imwrite(output_filename, output_image)
    if os.path.exists(output_filename):
        print("Image {} created!".format(output_filename))
    else:
        print("Image {} was not created. Check your permissions.".format(output_filename))